import asyncio
import aiohttp
import json
import pandas as pd
import time
import argparse
from typing import List, Dict
import random
import os
import sys

# 添加项目根目录到系统路径
sys.path.append(
    os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
)

# 导入letter模块的配置
from letter.config import current_config

# 根据环境配置构建默认API地址
DEFAULT_API_URL = (
    f"http://{current_config['host']}:{current_config['port']}/api/getLetters"
)

# 配置参数
parser = argparse.ArgumentParser(description="测试服务器并发处理能力")
parser.add_argument("--num_requests", type=int, default=10, help="并发请求数量")
parser.add_argument(
    "--excel_file", type=str, default="letter.xlsx", help="Excel文件路径"
)
parser.add_argument(
    "--server_url",
    type=str,
    default=DEFAULT_API_URL,
    help="服务器URL",
)
args = parser.parse_args()


async def send_request(session: aiohttp.ClientSession, data: Dict) -> Dict:
    """发送单个请求"""
    request_start_time = time.time()
    try:
        async with session.post(args.server_url, json=data) as response:
            result = await response.json()
            request_end_time = time.time()
            request_time = request_end_time - request_start_time
            
            # 添加请求时间信息到结果中
            result["request_time"] = request_time
            result["request_start_time"] = request_start_time
            result["request_end_time"] = request_end_time
            result["request_data"] = data  # 保存请求数据用于标识
            
            return result
    except Exception as e:
        request_end_time = time.time()
        request_time = request_end_time - request_start_time
        
        error_result = {
            "success": "false", 
            "message": str(e), 
            "data": {},
            "request_time": request_time,
            "request_start_time": request_start_time,
            "request_end_time": request_end_time,
            "request_data": data
        }
        return error_result


async def main():
    # 读取Excel文件
    try:
        df = pd.read_excel(args.excel_file, sheet_name="Sheet1")
    except Exception as e:
        print(f"读取Excel文件失败: {str(e)}")
        return

    # 准备测试数据
    test_data_list = []
    for _, row in df.iterrows():
        test_data = {
            "bh": str(row["信访举报件编号"]),
            "caseIntruction": str(row["原件描述"]),
            "txm": (
                str(row["条形码"])
                if "条形码" in row and not pd.isna(row["条形码"])
                else ""
            ),
        }
        test_data_list.append(test_data)

    # 如果测试数据不足，则重复使用现有数据
    while len(test_data_list) < args.num_requests:
        test_data_list.extend(test_data_list[: args.num_requests - len(test_data_list)])

    # 随机选择指定数量的测试数据
    selected_data = random.sample(test_data_list, args.num_requests)

    print(f"开始并发测试，请求数量: {args.num_requests}")
    start_time = time.time()

    # 创建会话并发送并发请求
    async with aiohttp.ClientSession() as session:
        tasks = [send_request(session, data) for data in selected_data]
        results = await asyncio.gather(*tasks)

    end_time = time.time()
    total_time = end_time - start_time

    # 统计结果
    success_count = sum(1 for r in results if r.get("success") == "true")
    fail_count = len(results) - success_count

    # 分析每个请求的处理时间
    request_times = [r.get("request_time", 0) for r in results if r.get("request_time")]
    if request_times:
        min_time = min(request_times)
        max_time = max(request_times)
        avg_time = sum(request_times) / len(request_times)
        median_time = sorted(request_times)[len(request_times) // 2]
    else:
        min_time = max_time = avg_time = median_time = 0

    print("\n测试结果统计:")
    print(f"总请求数: {len(results)}")
    print(f"成功请求数: {success_count}")
    print(f"失败请求数: {fail_count}")
    print(f"总耗时: {total_time:.2f}秒")
    print(f"平均每个请求耗时: {total_time/len(results):.2f}秒")
    print(f"QPS (每秒查询数): {len(results)/total_time:.2f}")
    
    print("\n单个请求时间分析:")
    print(f"最快请求: {min_time:.3f}秒")
    print(f"最慢请求: {max_time:.3f}秒")
    print(f"平均请求时间: {avg_time:.3f}秒")
    print(f"中位数请求时间: {median_time:.3f}秒")
    
    print("\n详细请求信息:")
    for i, result in enumerate(results, 1):
        status = "成功" if result.get("success") == "true" else "失败"
        request_time = result.get("request_time", 0)
        bh = result.get("request_data", {}).get("bh", "未知")
        print(f"  请求{i}: 编号{bh} - {status} - 耗时{request_time:.3f}秒")

    # 保存详细结果到文件
    result_file = f"concurrent_test_result_{int(time.time())}.json"
    with open(result_file, "w", encoding="utf-8") as f:
        json.dump(
            {
                "test_config": {
                    "num_requests": args.num_requests,
                    "server_url": args.server_url,
                    "excel_file": args.excel_file,
                },
                "summary": {
                    "total_requests": len(results),
                    "success_count": success_count,
                    "fail_count": fail_count,
                    "total_time": total_time,
                    "avg_time_per_request": total_time / len(results),
                    "qps": len(results) / total_time,
                    "individual_request_stats": {
                        "min_time": min_time,
                        "max_time": max_time,
                        "avg_time": avg_time,
                        "median_time": median_time
                    }
                },
                "detailed_results": results,
            },
            f,
            ensure_ascii=False,
            indent=2,
        )
    print(f"\n详细结果已保存到: {result_file}")


if __name__ == "__main__":
    asyncio.run(main())
